{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e6ac660a-892d-4aae-a936-e65e2b999c28",
   "metadata": {
    "papermill": {
     "duration": 0,
     "end_time": "2023-07-14T18:17:57.423606",
     "exception": false,
     "start_time": "2023-07-14T18:17:57.423606",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Segment-Anything"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f223aaf-3a0a-4f71-a6e1-4b3de62fe4f6",
   "metadata": {
    "papermill": {
     "duration": 0.008584,
     "end_time": "2023-07-14T18:17:57.440154",
     "exception": false,
     "start_time": "2023-07-14T18:17:57.431570",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Select objects specified by input points or box or both using Segmenting Anything Library.\n",
    "https://github.com/facebookresearch/segment-anything\n",
    "Output numpy datafiles in specified output directory.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d1936b3c-0e07-45d1-85df-4c699aa13a08",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-14T18:17:57.457508Z",
     "iopub.status.busy": "2023-07-14T18:17:57.457508Z",
     "iopub.status.idle": "2023-07-14T18:18:18.327552Z",
     "shell.execute_reply": "2023-07-14T18:18:18.327552Z"
    },
    "papermill": {
     "duration": 20.879081,
     "end_time": "2023-07-14T18:18:18.327552",
     "exception": false,
     "start_time": "2023-07-14T18:17:57.448471",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/facebookresearch/segment-anything.git\n",
      "  Cloning https://github.com/facebookresearch/segment-anything.git to c:\\users\\colin\\appdata\\local\\temp\\pip-req-build-nlkrrkpc\n",
      "  Resolved https://github.com/facebookresearch/segment-anything.git to commit 6fdee8f2727f4506cfbbe553e23b895e27956588\n",
      "  Preparing metadata (setup.py): started\n",
      "  Preparing metadata (setup.py): finished with status 'done'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  Running command git clone --filter=blob:none --quiet https://github.com/facebookresearch/segment-anything.git 'C:\\Users\\colin\\AppData\\Local\\Temp\\pip-req-build-nlkrrkpc'\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: opencv-python in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (4.8.0.74)\n",
      "Requirement already satisfied: numpy>=1.21.2 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from opencv-python) (1.25.0)\n",
      "Requirement already satisfied: torch in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (2.0.1)\n",
      "Requirement already satisfied: filelock in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch) (3.12.2)\n",
      "Requirement already satisfied: typing-extensions in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch) (4.6.3)\n",
      "Requirement already satisfied: sympy in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch) (1.12)\n",
      "Requirement already satisfied: networkx in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch) (3.1)\n",
      "Requirement already satisfied: jinja2 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch) (3.1.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from jinja2->torch) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from sympy->torch) (1.3.0)\n",
      "Requirement already satisfied: torchvision in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (0.15.2)\n",
      "Requirement already satisfied: numpy in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torchvision) (1.25.0)\n",
      "Requirement already satisfied: requests in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torchvision) (2.28.1)\n",
      "Requirement already satisfied: torch==2.0.1 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torchvision) (2.0.1)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torchvision) (10.0.0)\n",
      "Requirement already satisfied: filelock in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch==2.0.1->torchvision) (3.12.2)\n",
      "Requirement already satisfied: typing-extensions in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch==2.0.1->torchvision) (4.6.3)\n",
      "Requirement already satisfied: sympy in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch==2.0.1->torchvision) (1.12)\n",
      "Requirement already satisfied: networkx in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch==2.0.1->torchvision) (3.1)\n",
      "Requirement already satisfied: jinja2 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from torch==2.0.1->torchvision) (3.1.2)\n",
      "Requirement already satisfied: charset-normalizer<3,>=2 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from requests->torchvision) (2.1.1)\n",
      "Requirement already satisfied: idna<4,>=2.5 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from requests->torchvision) (3.4)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from requests->torchvision) (1.26.16)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from requests->torchvision) (2023.5.7)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from jinja2->torch==2.0.1->torchvision) (2.1.3)\n",
      "Requirement already satisfied: mpmath>=0.19 in d:\\python_projects\\anaconda\\envs\\claimed\\lib\\site-packages (from sympy->torch==2.0.1->torchvision) (1.3.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install git+https://github.com/facebookresearch/segment-anything.git\n",
    "!pip install opencv-python\n",
    "!pip install torch\n",
    "!pip install torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "300597ed-b9dc-4fdb-9cd4-da793c1aa5c5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-14T18:18:18.336040Z",
     "iopub.status.busy": "2023-07-14T18:18:18.336040Z",
     "iopub.status.idle": "2023-07-14T18:18:21.313988Z",
     "shell.execute_reply": "2023-07-14T18:18:21.313988Z"
    },
    "papermill": {
     "duration": 2.986436,
     "end_time": "2023-07-14T18:18:21.313988",
     "exception": false,
     "start_time": "2023-07-14T18:18:18.327552",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "import sys\n",
    "import re\n",
    "import cv2\n",
    "import numpy as np\n",
    "from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "29c527c9-ac90-4d98-a152-2b2301588aad",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-14T18:18:21.322004Z",
     "iopub.status.busy": "2023-07-14T18:18:21.322004Z",
     "iopub.status.idle": "2023-07-14T18:18:21.333862Z",
     "shell.execute_reply": "2023-07-14T18:18:21.333862Z"
    },
    "papermill": {
     "duration": 0.019874,
     "end_time": "2023-07-14T18:18:21.333862",
     "exception": false,
     "start_time": "2023-07-14T18:18:21.313988",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# model type\n",
    "model_type = os.environ.get(\"model_type\", \"vit_h\")  # default, vit_h, vit_l or vit_b\n",
    "\n",
    "# checkpoint\n",
    "# different model type requires different check point\n",
    "checkpoint_path = os.environ.get(\"checkpoint_path\")\n",
    "\n",
    "# input image\n",
    "input_image_path = os.environ.get(\"input_image_path\")\n",
    "\n",
    "# input points\n",
    "# input either input points or box or both to specify the selected object\n",
    "# format:\n",
    "# x,y,x,y,x,y... for points. Multiple allowed\n",
    "input_points = os.environ.get(\"input_array\", \"None\")  # e.g. 100,200\n",
    "\n",
    "# input box\n",
    "# format:\n",
    "# x,y,x,y for box. Only one box allowed\n",
    "input_box = os.environ.get(\"input_box\", \"None\")  # e.g. 100,200,300,400\n",
    "\n",
    "# temporal data storage for local execution\n",
    "data_dir = os.environ.get('data_dir', '../../data/')\n",
    "\n",
    "# dummy_output (to be fixed once C3 supports < 1 outputs)\n",
    "# output_dummy = os.environ.get('output_dummy', 'output_dummy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "81032697-8712-4f91-99d5-66351a4341d3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-14T18:18:21.333862Z",
     "iopub.status.busy": "2023-07-14T18:18:21.333862Z",
     "iopub.status.idle": "2023-07-14T18:18:21.350968Z",
     "shell.execute_reply": "2023-07-14T18:18:21.350968Z"
    },
    "papermill": {
     "duration": 0.017106,
     "end_time": "2023-07-14T18:18:21.350968",
     "exception": false,
     "start_time": "2023-07-14T18:18:21.333862",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "parameters = list(\n",
    "    map(lambda s: re.sub('$', '\"', s),\n",
    "        map(\n",
    "            lambda s: s.replace('=', '=\"'),\n",
    "            filter(\n",
    "                lambda s: s.find('=') > -1 and bool(re.match(r'[A-Za-z0-9_]*=[.\\/A-Za-z0-9]*', s)),\n",
    "                sys.argv\n",
    "            )\n",
    "    )))\n",
    "\n",
    "for parameter in parameters:\n",
    "    logging.warning('Parameter: ' + parameter)\n",
    "    exec(parameter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "df9f546d-a60e-47b7-b413-a302077dcb92",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-07-14T18:18:21.359013Z",
     "iopub.status.busy": "2023-07-14T18:18:21.359013Z",
     "iopub.status.idle": "2023-07-14T18:20:15.031375Z",
     "shell.execute_reply": "2023-07-14T18:20:15.031375Z"
    },
    "papermill": {
     "duration": 113.680407,
     "end_time": "2023-07-14T18:20:15.031375",
     "exception": false,
     "start_time": "2023-07-14T18:18:21.350968",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "both\n",
      "done\n"
     ]
    }
   ],
   "source": [
    "# get masks from a given prompt\n",
    "device = \"cuda\"\n",
    "sam = sam_model_registry[model_type](checkpoint=checkpoint_path)\n",
    "predictor = SamPredictor(sam)\n",
    "input_image = cv2.imread(input_image_path)\n",
    "predictor.set_image(input_image)\n",
    "\n",
    "# convert input points and box into array\n",
    "if input_points != \"None\":\n",
    "    input_points_list = input_points.split(',')\n",
    "    points = [[int(input_points_list[i]), int(input_points_list[i+1])] for i in range(0, len(input_points_list), 2)]\n",
    "if input_box != \"None\":\n",
    "    input_box_list = input_box.split(',')\n",
    "    box = [int(x) for x in input_box_list]\n",
    "\n",
    "# process output with different inputs mix\n",
    "# obly points\n",
    "if input_points != \"None\" and input_box == \"None\":\n",
    "    print(\"points\")\n",
    "    input_label = np.array([1] * len(points))\n",
    "    points = np.array(points)\n",
    "    masks, scores, logits = predictor.predict(\n",
    "        point_coords=points,\n",
    "        point_labels=input_label,\n",
    "        multimask_output=True,\n",
    "    )\n",
    "\n",
    "# only box\n",
    "elif input_points == \"None\" and input_box != \"None\":\n",
    "    box = np.array(box)\n",
    "    masks, scores, logits = predictor.predict(\n",
    "        point_coords=None,\n",
    "        point_labels=None,\n",
    "        box=box[None, :],\n",
    "        multimask_output=False,\n",
    "    )\n",
    "# both point and box\n",
    "elif input_points != \"None\" and input_box != \"None\":\n",
    "    input_label = np.array([1] * len(points))\n",
    "    points = np.array(points)\n",
    "    box = np.array(box)\n",
    "    masks, scores, logits = predictor.predict(\n",
    "        point_coords=points,\n",
    "        point_labels=input_label,\n",
    "        box=box,\n",
    "        multimask_output=False,\n",
    "    )\n",
    "\n",
    "# save result into given directory\n",
    "output_directory, output_filename = os.path.split(data_dir)\n",
    "os.makedirs(output_directory, exist_ok=True)\n",
    "np.save(data_dir + \"get_masks\", masks)\n",
    "np.save(data_dir + \"scores\", scores)\n",
    "np.save(data_dir + \"logits\", logits)\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 141.727187,
   "end_time": "2023-07-14T18:20:15.941273",
   "environment_variables": {},
   "exception": null,
   "input_path": "D:\\Python_projects\\Claimed\\component-library\\component-library\\segment-anything\\get-masks.ipynb",
   "output_path": "D:\\Python_projects\\Claimed\\component-library\\component-library\\segment-anything\\get-masks.ipynb",
   "parameters": {},
   "start_time": "2023-07-14T18:17:54.214086",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
